import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers, activations, optimizers, losses
from python_ai.common.xcommon import *


def ConvBnRelu(filters, size, strides=(1, 1), name='NNm', padding='same'):
    conv_name = name + '_conv_' + repr2valid_name(size) + '_' + repr2valid_name(strides) + '_' + repr2valid_name(padding)
    conv_name = shrink_underscore(conv_name)
    model = keras.Sequential([
        layers.Conv2D(
            filters,
            size,
            strides=strides,
            padding=padding,
            # use_bias=False,
            name=conv_name
        ),
        layers.BatchNormalization(name=name + '_bn', scale=False),
        layers.ReLU(name=name + '_rl')
    ], name=name)
    return model


class InceptionBlock(keras.Model):

    def __init__(self, branch_out_ch, shrink=False, name='NNm', **kwargs):
        super().__init__(name=name, **kwargs)

        if shrink:
            strides = (2, 2)
        else:
            strides = (1, 1)
        self.branch01 = ConvBnRelu(branch_out_ch, (1, 1), strides, name + '_b01')
        self.branch02_1 = ConvBnRelu(branch_out_ch, (1, 1), strides, name + '_b02_1')
        self.branch02_2 = ConvBnRelu(branch_out_ch, (3, 3), (1, 1), name + '_b02_2')

        self.branch03_1 = ConvBnRelu(branch_out_ch, (1, 1), strides, name + '_b03_1')
        self.branch03_2 = ConvBnRelu(branch_out_ch, (3, 3), (1, 1), name + '_b03_2')
        self.branch03_3 = ConvBnRelu(branch_out_ch, (3, 3), (1, 1), name + '_b03_3')

        self.branch04_1 = layers.AvgPool2D((3, 3), (1, 1), padding='same', name=name + '_b04_1')
        self.branch04_2 = ConvBnRelu(branch_out_ch, (1, 1), strides, name + '_b04_2')

    def call(self, x, training=None):
        x1 = self.branch01(x, training=training)
        
        x2 = self.branch02_1(x, training=training)
        x2 = self.branch02_2(x2, training=training)

        x3 = self.branch03_1(x, training=training)
        x3 = self.branch03_2(x3, training=training)
        x3 = self.branch03_3(x3, training=training)

        x4 = self.branch04_1(x, training=training)
        x4 = self.branch04_2(x4, training=training)

        res = tf.concat((x1, x2, x3, x4), axis=3)
        return res


class InceptionNet(keras.Model):

    def __init__(self, n_cls, init_ch, n_blocks, activation='softmax', **kwargs):
        super().__init__(**kwargs)

        self.conv = ConvBnRelu(init_ch, (3, 3), (1, 1), '01')
        self.inception_layers = []
        ch = init_ch
        for block_id in range(n_blocks):
            for layer_id in range(2):
                if layer_id == 0:
                    ch *= 2
                    shrink = True
                else:
                    shrink = False
                layer = InceptionBlock(ch, shrink=shrink, name=f'B{block_id + 1}{layer_id + 1}')
                self.inception_layers.append(layer)
        self.g_avg_pool = layers.GlobalAveragePooling2D(name='GlbAvgPl')
        self.fc = layers.Dense(n_cls, activation=activation, name='FC')

    def call(self, x, training=None):
        x = self.conv(x, training=training)
        for layer in self.inception_layers:
            x = layer(x, training=training)
        x = self.g_avg_pool(x, training=training)
        x = self.fc(x, training=training)
        return x


if '__main__' == __name__:

    tf.random.set_seed(777)

    (x_train,y_train),(x_test,y_test)=tf.keras.datasets.cifar10.load_data()

    x_train=x_train.reshape([-1,32,32,3]).astype('float32') / 255
    x_test=x_test.reshape([-1,32,32,3]).astype('float32') / 255

    model = InceptionNet(10, 16, 2)
    model.build(input_shape=(None, 32, 32, 3))
    model.summary(line_length=120)

    model.compile(loss=losses.SparseCategoricalCrossentropy(),
                  optimizer=optimizers.Adam(lr=0.001),
                  metrics=['accuracy'])

    history=model.fit(x_train, y_train, batch_size=64, epochs=3, validation_split=0.1)

    score=model.evaluate(x_test, y_test, batch_size=64)
    print('accuracy',score[1])
    print('loss',score[0])
